from torchvision import transforms
from torchvision.datasets import CIFAR10
from torch.nn import Linear, Module, Flatten, MaxPool2d, Conv2d, Sequential, CrossEntropyLoss
import torch

dataset = CIFAR10(root="./dataset", train=False, transform=transforms.ToTensor())
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)

class MyCifarModel(Module):
    def __init__(self):
        super(MyCifarModel, self).__init__()
        self.model1 = Sequential(
            Conv2d(3, 32, 3, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, x):
        return self.model1(x)

my_model = MyCifarModel()
cel = CrossEntropyLoss()
optim = torch.optim.SGD(my_model.parameters(), lr=0.01)

for epoch in range(20):
    total_loss = 0.0
    for data in dataloader:
        imgs, labels = data
        outputs = my_model(imgs)
        loss = cel(outputs, labels)

        optim.zero_grad()
        # 使用损失函数进行反向传播
        loss.backward()
        optim.step()
        total_loss += loss.item()
    print(total_loss)